# 数据库操作工具

import pymysql
import sqlite3
from config.config import get_config


class DBHelper:
    """数据库操作辅助类，支持MySQL和SQLite"""
    
    def __init__(self):
        """初始化数据库连接"""
        self.config = get_config()
        self.db_type = self.config.get('database', 'type')
        self.connection = None
        self.cursor = None
    
    def connect(self):
        """连接到数据库"""
        if self.db_type.lower() == 'mysql':
            self.connection = pymysql.connect(
                host=self.config.get('database', 'host'),
                port=int(self.config.get('database', 'port')),
                user=self.config.get('database', 'user'),
                password=self.config.get('database', 'password'),
                database=self.config.get('database', 'database'),
                charset='utf8mb4'
            )
        elif self.db_type.lower() == 'sqlite':
            db_path = self.config.get('database', 'path')
            self.connection = sqlite3.connect(db_path)
        else:
            raise ValueError(f"不支持的数据库类型: {self.db_type}")
        
        self.cursor = self.connection.cursor()
        return self.connection
    
    def execute_query(self, sql, params=None):
        """执行查询SQL"""
        try:
            if not self.connection:
                self.connect()
            
            self.cursor.execute(sql, params or ())
            return self.cursor.fetchall()
        except Exception as e:
            print(f"查询执行错误: {e}")
            raise
    
    def execute_update(self, sql, params=None):
        """执行更新SQL（INSERT, UPDATE, DELETE）"""
        try:
            if not self.connection:
                self.connect()
            
            self.cursor.execute(sql, params or ())
            self.connection.commit()
            return self.cursor.rowcount
        except Exception as e:
            self.connection.rollback()
            print(f"更新执行错误: {e}")
            raise
    
    def close(self):
        """关闭数据库连接"""
        if self.cursor:
            self.cursor.close()
        if self.connection:
            self.connection.close()
    
    def __enter__(self):
        """支持with语句"""
        self.connect()
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        """退出with语句块时自动关闭连接"""
        self.close()


# 常用数据库操作方法
def get_user_by_username(username):
    """根据用户名获取用户信息"""
    with DBHelper() as db:
        sql = "SELECT id, username, password FROM tb_users WHERE username = %s"
        result = db.execute_query(sql, (username,))
        return result[0] if result else None


def get_product_by_id(product_id):
    """根据ID获取商品信息"""
    with DBHelper() as db:
        sql = "SELECT id, name, price FROM tb_goods WHERE id = %s"
        result = db.execute_query(sql, (product_id,))
        return result[0] if result else None


def get_order_by_id(order_id):
    """根据ID获取订单信息"""
    with DBHelper() as db:
        sql = "SELECT id, user_id, total_amount, status FROM tb_order_info WHERE id = %s"
        result = db.execute_query(sql, (order_id,))
        return result[0] if result else None


def clean_test_data():
    """清理测试数据，用于测试后的数据恢复"""
    with DBHelper() as db:
        # 删除测试用户创建的订单
        db.execute_update("DELETE FROM tb_order_info WHERE user_id IN (SELECT id FROM tb_users WHERE username LIKE 'test_%')")
        # 删除测试用户的购物车
        db.execute_update("DELETE FROM tb_cart WHERE user_id IN (SELECT id FROM tb_users WHERE username LIKE 'test_%')")
        # 删除测试用户创建的地址
        db.execute_update("DELETE FROM tb_address WHERE user_id IN (SELECT id FROM tb_users WHERE username LIKE 'test_%')")
        # 删除测试用户
        db.execute_update("DELETE FROM tb_users WHERE username LIKE 'test_%'")